import sys
import json
import logging
import asyncio
from typing import Optional, Dict, Any, List
from contextlib import AsyncExitStack
from mcp.client.session import ClientSession
from mcp.client.stdio import StdioServerParameters
from mcp.client.stdio import stdio_client
from mcp.client.sse import sse_client

logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
)


class MCPTool:
    """ 封装 MCP 返回的工具信息 """

    def __init__(self, name: str, description: str, input_schema: Dict[str, Any]) -> None:
        self.name: str = name
        self.description: str = description
        self.input_schema: Dict[str, Any] = input_schema


class MCPClient:
    def __init__(self):
        self._session_context = None
        self._streams_context = None
        self._cleanup_lock = asyncio.Lock()
        self.exit_stack = AsyncExitStack()
        self.session: Optional[ClientSession] = None

    async def connect_to_server(self, server_script_path: str = None, service_url: str = None, env: dict = None):
        """
        连接到MCP服务器

        Args:
            server_script_path: Path to the server script (.py or .js)
            service_url: URL of the service
            env: env

        Author:
            zero
        """
        if service_url is not None:
            await self._connect_to_sse_server(service_url)
        elif server_script_path is not None:
            await self._connect_to_stdio_server(server_script_path, env=env)
        else:
            raise ValueError("Either server_script_path or service_url must be provided.")

    async def _connect_to_sse_server(self, service_url: str):
        try:
            self._streams_context = sse_client(url=service_url)
            streams = await self._streams_context.__aenter__()
            self._session_context = ClientSession(*streams)
            self.session: ClientSession = await self._session_context.__aenter__()
            await self.session.initialize()
        except Exception as e:
            await self.cleanup()
            if len(e.args) >= 2:
                if "401 Unauthorized" in str(e.args[1]):
                    raise Exception("SSE Client error 401 Unauthorized")
                if "Request URL is missing" in str(e.args[1]):
                    raise Exception("SSE Client error 404 Request URL is missing")
                if "ConnectError" in str(e.args[1]):
                    raise Exception("SSE Client error 404 Not known")
            raise Exception("SSE Client error " + str(e))

    async def _connect_to_stdio_server(self, server_script_path: str, env: dict = None):
        is_python = server_script_path.endswith(".py")
        is_js = server_script_path.endswith(".js")
        if not (is_python or is_js):
            raise ValueError("Server script must be a .py or .js file")

        command = sys.executable if is_python else "node"
        server_params = StdioServerParameters(
            command=command, args=[server_script_path], env=env
        )

        stdio_transport = await self.exit_stack.enter_async_context(
            stdio_client(server_params)
        )

        self.stdio, self.write = stdio_transport
        self.session = await self.exit_stack.enter_async_context(
            ClientSession(self.stdio, self.write)
        )
        await self.session.initialize()

    async def get_tools_dict(self) -> List[Dict[str, Any]]:
        """
        获取工具字典

        Returns:
            List[Dict[str, Any]]

        Author:
            zero
        """
        _tools_list = []
        response = await self.session.list_tools()
        for tool in response.tools:
            _tools_list.append({
                "name": tool.name,
                "is_enable": True,
                "description": tool.description,
                "inputSchema": tool.inputSchema
            })
        return _tools_list

    async def get_tools_list(self) -> List[MCPTool]:
        """
        获取工具列表

        Returns:
            List[MCPTool]

        Author:
            zero
        """
        tools_response = await self.session.list_tools()
        tools = []
        for item in tools_response:
            if isinstance(item, tuple) and item[0] == "tools":
                for tool in item[1]:
                    tools.append(MCPTool(tool.name, tool.description, tool.inputSchema))
        return tools

    async def call_tool(self, tool_name: str, arguments: Dict[str, Any], retries: int = 2, delay: float = 1.0):
        """
        执行指定工具

        Args:
            tool_name (str): 工具名称
            arguments (Dict[str, Any]): 工具参数
            retries (int): 重试次数
            delay (float): 重试间隔秒数

        Returns:
            any: 工具调用结果

        Author:
            zero
        """
        attempt = 0
        while attempt < retries:
            try:
                return await self.session.call_tool(name=tool_name, arguments=arguments)
            except Exception as e:
                attempt += 1
                logging.warning(f"Error executing tool: {e}. Attempt {attempt} of {retries}.")
                if attempt < retries:
                    logging.info(f"Retrying in {delay} seconds...")
                    await asyncio.sleep(delay)
                else:
                    logging.error("Max retries reached. Failing.")
                    raise
            finally:
                await self.cleanup()
        return None

    async def cleanup(self):
        """ Clean up resources """
        async with self._cleanup_lock:
            try:
                # 先清理 session
                if self._session_context:
                    try:
                        await self._session_context.__aexit__(None, None, None)
                    except Exception as e:
                        logging.error(f"Error closing session: {e}")
                    finally:
                        self._session_context = None

                # 然后清理 streams
                if self._streams_context:
                    try:
                        await self._streams_context.__aexit__(None, None, None)
                    except Exception as e:
                        logging.error(f"Error closing streams: {e}")
                    finally:
                        self._streams_context = None

                # 最后清理 exit_stack
                if self.exit_stack:
                    try:
                        await self.exit_stack.aclose()
                    except Exception as e:
                        logging.error(f"Error closing exit stack: {e}")
                    finally:
                        self.exit_stack = AsyncExitStack()

                self.session = None
            except Exception as e:
                logging.error(f"Error during cleanup of server: {e}")


class MCPClientPool:
    def __init__(self):
        self.servers = {}
        self.tools_all = []
        self.tools_by_server = {}
        self.exit_stack = AsyncExitStack()

    async def connect_to_servers(self, servers_config: Dict[str, Any]) -> None:
        mcp_servers = servers_config.get("mcpServers", {})
        for server_name, srv_config in mcp_servers.items():
            # 链接服务器
            server = None
            if srv_config.get("url"):
                server = MCPClient()
                await server.connect_to_server(service_url=srv_config.get("url"))

            # 工具集获取
            tools = srv_config.get("tools", [])
            if not srv_config.get("tools"):
                tools = await server.get_tools_list()
            else:
                tools = [MCPTool(
                            name=t["name"],
                            description=t["description"],
                            input_schema=t["inputSchema"]) for t in tools]

            # 字典存服务
            self.servers[server_name] = server
            self.tools_by_server[server_name] = tools

            # 工具重命名
            for tool in tools:
                function_name = f"{server_name}_{tool.name}"
                self.tools_all.append({
                    "type": "function",
                    "function": {
                        "name": function_name,
                        "description": tool.description,
                        "input_schema": tool.input_schema
                    }
                })

            # 转换为 (Function Calling)
            self.tools_all = await self.transform_json(self.tools_all)

    @classmethod
    async def transform_json(cls, json_data: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """ 将工具的 input_schema 转换为 OpenAI 所需的 parameters 格式，并删除多余字段 """
        result = []
        for item in json_data:
            if not isinstance(item, dict) or "type" not in item or "function" not in item:
                continue

            old_func = item["function"]
            if not isinstance(old_func, dict) or "name" not in old_func or "description" not in old_func:
                continue

            new_func = {
                "name": old_func["name"],
                "description": old_func["description"],
                "parameters": {}
            }

            if "input_schema" in old_func and isinstance(old_func["input_schema"], dict):
                old_schema = old_func["input_schema"]
                new_func["parameters"]["type"] = old_schema.get("type", "object")
                new_func["parameters"]["properties"] = old_schema.get("properties", {})
                new_func["parameters"]["required"] = old_schema.get("required", [])

            result.append({
                "type": item["type"],
                "function": new_func
            })
        return result

    async def chat_loop(self, llmServer, messages: list):
        """ 多服务器MCP, 主动循环对话 """
        try:
            tools = self.tools_all
            while True:
                calls = None
                stream = await llmServer.llm(messages=messages, tools=tools, stream=True)
                async for chunk in stream:
                    choice = chunk.choices[0]
                    if choice.delta.tool_calls is not None:
                        for i, t in enumerate(choice.delta.tool_calls):
                            if calls:
                                arg = t.function.arguments or ""
                                calls.choices[0].delta.tool_calls[i].function.arguments += arg
                            else:
                                calls = chunk
                    elif choice.finish_reason == "tool_calls":
                        calls.choices[0].finish_reason = "tool_calls"
                        yield calls
                    else:
                        yield chunk

                if calls:
                    messages, resp = await self._resp_messages(messages, calls)
                    calls.choices[0].delta.content = resp
                    yield calls
                    continue

                break
        finally:
            await self.cleanup()

    async def _resp_messages(self, messages, response):
        """
        将模型返回的工具调用解析执行, 并将结果追加到消息队列中
        """
        response_results = {}
        function_call_messages = response.choices[0].delta.tool_calls
        messages.append(response.choices[0].delta.model_dump())
        for function_call_message in function_call_messages:
            tool_name = function_call_message.function.name
            tool_args = json.loads(function_call_message.function.arguments)
            function_response = await self._call_mcp_tool(tool_name, tool_args)
            response_results[tool_name] = function_response
            messages.append({
                "role": "tool",
                "content": function_response,
                "tool_call_id": function_call_message.id,
            })
        return messages, json.dumps(response_results, ensure_ascii=False)

    async def _call_mcp_tool(self, tool_full_name: str, tool_args: Dict[str, Any]):
        """
        根据 "serverName_toolName" 格式调用相应 MCP 工具
        """
        parts = tool_full_name.split("_", 1)
        if len(parts) != 2:
            return f"无效的工具名称: {tool_full_name}"

        server_name, tool_name = parts
        server = self.servers.get(server_name)
        if not server:
            return f"找不到服务器: {server_name}"

        resp = await server.call_tool(tool_name, tool_args)
        if resp.content:
            content = []
            for item in resp.content:
                content.append(item.model_dump())
            return json.dumps(content, ensure_ascii=False)
        return "工具执行无输出"

    async def cleanup(self):
        """ Clean up resources """
        for key, srv in self.servers.items():
            await srv.cleanup()
        if self.exit_stack:
            await self.exit_stack.aclose()
        self.servers = {}
        self.tools_all = []
        self.tools_by_server = {}
